package com.sofun.ai.config;

import com.sofun.ai.StudyAiApplication;
import com.sofun.ai.assistant.RagAssistant;
import com.sofun.ai.assistant.AiAssistant;
import com.sofun.ai.tools.Tools;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Collection;
import java.util.List;
import java.util.Objects;

@Configuration
public class AiConfig {

/*    @Resource
    RedisEmbeddingStore redisEmbeddingStore;*/

    /**
     * 自定义QwenChatModel名字为qwenChatModel1（避免与百炼的qwenChatModel名字重复，保证下面加载的是自定义的qwenChatModel1），主要是注入监听，打印日志
     * */
    @Bean
    @Primary
    QwenChatModel qwenChatModel1(QwenProperties properties, List<ChatModelListener> chatModelListenerList) {
        return QwenChatModel.builder()
                .apiKey(properties.getApiKey())
                .modelName(properties.getModelName())
                .listeners(chatModelListenerList)
                .build();
    }

    /**
     * @param qwenChatModel1 自定义的ChatLanguageModel，可使用框架的ChatLanguageModel对象
     * @param qwenStreamingChatModel 流式输出的大模型对象
     * @param tools function call工具集合
     * @param retrievalAugmentor 检索增强器（自定义，可以结合embedding和web search输出）
     * */
    @Bean
    @ConditionalOnClass(QwenChatModel.class)
    public RagAssistant ragAssistant(ChatLanguageModel qwenChatModel1,
                                     StreamingChatLanguageModel qwenStreamingChatModel,
                                     Collection<Tools> tools,
                                     ChatMemoryProvider chatMemoryProvider,
                                     RetrievalAugmentor retrievalAugmentor
    ) {
        // 对话记忆(与chatMemoryProvider选择一个使用)
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        // 为Assistant动态代理对象  chat  --->  对话内容存储ChatMemory----> 聊天记录ChatMemory取出来 ---->放入到当前对话中
        return AiServices.builder(RagAssistant.class)
                .tools(tools.toArray())//注意这个地方传集合的话必须传Collection<Object>，不能传List<Object>类型的集合
               // .chatMemory(chatMemory)
                .chatLanguageModel(qwenChatModel1)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .retrievalAugmentor(retrievalAugmentor)
                .chatMemoryProvider(chatMemoryProvider)
                .build();
    }


    /**定义与大模型对话的Assistant
     * @param qwenChatModel1 自定义的ChatLanguageModel，可使用框架的ChatLanguageModel对象
     * @param qwenStreamingChatModel 流式输出的大模型对象
     * @param tools function call工具集合
     * */
    @Bean
    public AiAssistant aiAssistant(ChatLanguageModel qwenChatModel1,
                                   StreamingChatLanguageModel qwenStreamingChatModel,
                                   ChatMemoryProvider chatMemoryProvider,
                                   Collection<Tools> tools
    ) {
        return AiServices.builder(AiAssistant.class)
                .chatLanguageModel(qwenChatModel1)
                .tools(tools.toArray())//注意这个地方传集合的话必须传Collection<Object>，不能传List<Object>类型的集合
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .build();
    }


    /**
     * ChatMemoryProvider 是一个工厂接口（Factory Interface），用于根据请求上下文动态创建或提供 ChatMemory 实例
     * */
    @Bean
    public ChatMemoryProvider chatMemoryProvider(){
        //自定义ChatMemoryProvider
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(10)
                .chatMemoryStore(dbChatMemoryStore())
                .build();
    }

    /**
     * 对话记忆存存储器
     * @return 对话记忆存存储器
     */
    @Bean
    public ChatMemoryStore dbChatMemoryStore() {
        return new DbChatMemoryStore();
    }

    /**
     *
     * @return 向量数据库(基于内存内存)
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 项目启动时提前加载document到向量数据库
     * */
    @Bean
    public CommandLineRunner ingestTermOfServiceToVectorStore(QwenEmbeddingModel qwenEmbeddingModel,EmbeddingStore embeddingStore) throws URISyntaxException {
        Path documentPath = Paths.get(Objects.requireNonNull(StudyAiApplication.class.getClassLoader().getResource("rag/terms-of-service.txt")).toURI());

        return arg-> {
            TextDocumentParser textDocumentParser = new TextDocumentParser();
            Document document = FileSystemDocumentLoader.loadDocument(documentPath, textDocumentParser);

            DocumentByRegexSplitter splitter = new DocumentByRegexSplitter(
                    "\\n\\d+\\.",  // 匹配 "1. 标题" 格式
                    "\n",            // 保留换行符作为段落连接符
                    100,             // 每个段最多 500 字符
                    20,               // 段间重叠 50 字符以保持连贯性
                    new DocumentByCharacterSplitter(100,50)
            );

            List<TextSegment> segments = splitter.split(document);

            // 向量化
            List<Embedding> embeddings = qwenEmbeddingModel.embedAll(segments).content();
            //分词后的数据给向量数据库
            embeddingStore.addAll(embeddings,segments);
        };
    }


}
